#include "threadsafe_queue.h"

#include <gtest/gtest.h>

#include <atomic>
#include <cstdint>
#include <mutex>
#include <thread>
#include <unordered_set>
#include <vector>

/**
 * 多生产者多消费者测试：
 * - 4 个生产者、4 个消费者；
 * - 每个生产者推送 1000 个唯一任务；
 * - 所有消费者收到 close() 通知后自动退出；
 * - 断言消费总数与期待相符且没有重复任务。
 */
TEST(ProducerConsumerTest, MultiProducerMultiConsumer) {
    ThreadSafeQueue<std::int64_t> queue;

    constexpr int kProducerCount = 4;
    constexpr int kConsumerCount = 4;
    constexpr std::int64_t kItemsPerProducer = 1000;
    constexpr std::int64_t kTotalItems =
        kProducerCount * kItemsPerProducer;

    std::atomic<std::int64_t> produced{0};
    std::atomic<std::int64_t> consumed{0};

    std::vector<std::thread> producers;
    std::vector<std::thread> consumers;

    // 使用集合验证是否存在重复消费，同步访问需要互斥锁
    std::unordered_set<std::int64_t> consumedItems;
    std::mutex consumedMutex;

    // 启动生产者线程
    for (int pid = 0; pid < kProducerCount; ++pid) {
        producers.emplace_back([pid, &queue, &produced]() {
            const std::int64_t base = pid * kItemsPerProducer;
            for (std::int64_t i = 0; i < kItemsPerProducer; ++i) {
                queue.push(base + i);  // 唯一任务编号
                ++produced;
            }
        });
    }

    // 启动消费者线程
    for (int cid = 0; cid < kConsumerCount; ++cid) {
        consumers.emplace_back([&queue, &consumed, &consumedItems,
                                &consumedMutex]() {
            std::int64_t value;
            while (queue.pop(value)) {
                {
                    std::lock_guard<std::mutex> lock(consumedMutex);
                    consumedItems.insert(value);
                }
                ++consumed;
            }
            // pop() 返回 false 表示队列关闭且为空 -> 正常退出
        });
    }

    // 等待所有生产者完成任务，然后关闭队列
    for (auto& t : producers) {
        t.join();
    }
    queue.close();  // 使用 notify_all() 唤醒所有等待的消费者

    // 等待所有消费者退出
    for (auto& t : consumers) {
        t.join();
    }

    EXPECT_EQ(produced.load(), kTotalItems);
    EXPECT_EQ(consumed.load(), kTotalItems);
    EXPECT_EQ(consumedItems.size(), static_cast<std::size_t>(kTotalItems));
}

/**
 * 验证：在队列已关闭的情况下 push 会抛出异常。
 */
TEST(ProducerConsumerTest, PushAfterCloseThrows) {
    ThreadSafeQueue<int> queue;
    queue.close();
    EXPECT_THROW(queue.push(42), std::runtime_error);
}

// 可选的自定义入口（如果不链接 gtest_main 库）
// 将所有测试进行初始化并运行。
int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}